#include "types.h"
#include "param.h"
#include "riscv.h"
#include "defs.h"
#include "spinlock.h"
#include "proc.h"
#include "printf.h"

extern volatile uint64 ticks;
extern struct spinlock tickslock;

// 便捷：sleep 指定 tick 数
static void
sleep_ticks(int n)
{
  acquire(&tickslock);
  uint64 start = ticks;
  while ((int)(ticks - start) < n) {
    sleep((void*)&ticks, &tickslock);
  }
  release(&tickslock);
}

// ---------------- 测试1：进程创建/销毁 ----------------

static void
simple_task(void)
{
  // 简单任务：打印一次后返回（返回后 kthread_entry 调用 exit_process(0)）
  struct proc *p = myproc();
  int pid = p ? p->pid : -1;
  printf("simple_task: pid=%d started\n", pid);
  // 可加一点忙等，放大调度窗口
  for (volatile int i = 0; i < 200000; i++) {}
  // 返回即可
}

static void
test_process_creation(void)
{
  printf("=== Test: process creation ===\n");

  int pid0 = create_process(simple_task);
  if (pid0 <= 0) {
    printf("create_process(simple_task) failed\n");
    return;
  }

  int count = 0;
  for (int i = 0; i < NPROC + 8; i++) {
    int pid = create_process(simple_task);
    if (pid > 0) {
      count++;
    } else {
      // 到达容量或失败
      break;
    }
  }
  printf("created %d extra processes\n", count);

  // 回收所有子进程
  int waited = 0;
  int st;
  while (1) {
    int w = wait_process(&st);
    if (w < 0) break;
    waited++;
  }
  printf("waited %d children, done\n", waited);
}

// ---------------- 测试2：调度器/时钟抢占 ----------------

static volatile int sched_run_flag = 1;

static void
cpu_intensive_task(void)
{
  int local = 0;

  // 忙等直到被主控置位停止
  while (sched_run_flag) {
    // 一些“工作”
    for (volatile int i = 0; i < 100000; i++) { }
    local++;
    if ((local & 0x3ff) == 0) {
      // 避免过度输出
      // printf("cpu_task: work=%d\n", local);
    }
  }
  // 返回后退出
}

static void
test_scheduler(void)
{
  printf("=== Test: scheduler (preempt) ===\n");

  sched_run_flag = 1;

  // 创建多个计算密集型任务
  int n = 3;
  for (int i = 0; i < n; i++) {
    int pid = create_process(cpu_intensive_task);
    if (pid < 0) {
      printf("create_process(cpu_task %d) failed\n", i);
    }
  }

  // 观察一段时间（约 1/2 秒）
  sleep_ticks(50);

  // 通知停止
  sched_run_flag = 0;

  // 回收所有子进程
  int reaped = 0;
  while (reaped < n) {
    if (wait_process(0) > 0)
      reaped++;
    else
      break;
  }
  printf("scheduler test reaped=%d\n", reaped);
}

// ---------------- 测试3：sleep/wakeup 同步（生产者-消费者） ----------------

#define RINGSZ 16
#define PROD_ITEMS 100

static struct {
  int buf[RINGSZ];
  int r;      // 读索引
  int w;      // 写索引
  int cnt;    // 当前元素个数
  struct spinlock lock;
  int produced;
  int consumed;
  // 用两个不同通道地址避免混淆
  char ch_not_empty;
  char ch_not_full;
} rb;

static void
rb_init(void)
{
  rb.r = rb.w = rb.cnt = 0;
  rb.produced = rb.consumed = 0;
  initlock(&rb.lock, "ring");
}

static void
producer_task(void)
{
  for (int i = 1; i <= PROD_ITEMS; i++) {
    acquire(&rb.lock);
    while (rb.cnt == RINGSZ) {
      // 满了，等 not_full
      sleep((void*)&rb.ch_not_full, &rb.lock);
    }
    rb.buf[rb.w] = i;
    rb.w = (rb.w + 1) % RINGSZ;
    rb.cnt++;
    rb.produced++;
    // 唤醒消费者
    wakeup((void*)&rb.ch_not_empty);
    release(&rb.lock);
  }
  // 返回结束
}

static void
consumer_task(void)
{
  int sum = 0;
  for (int i = 0; i < PROD_ITEMS; i++) {
    acquire(&rb.lock);
    while (rb.cnt == 0) {
      // 空了，等 not_empty
      sleep((void*)&rb.ch_not_empty, &rb.lock);
    }
    int v = rb.buf[rb.r];
    rb.r = (rb.r + 1) % RINGSZ;
    rb.cnt--;
    rb.consumed++;
    // 唤醒生产者
    wakeup((void*)&rb.ch_not_full);
    release(&rb.lock);

    sum += v;
  }
  printf("consumer: sum=%d\n", sum);
}

static void
test_synchronization(void)
{
  printf("=== Test: synchronization (pc) ===\n");
  rb_init();

  int ppid = create_process(producer_task);
  int cpid = create_process(consumer_task);
  if (ppid < 0 || cpid < 0) {
    printf("create producer/consumer failed\n");
    return;
  }

  // 等待两者退出
  int got = 0;
  while (got < 2) {
    if (wait_process(0) > 0) got++;
    else break;
  }

  // 校验
  int ok = (rb.produced == PROD_ITEMS) && (rb.consumed == PROD_ITEMS);
  printf("pc result: produced=%d consumed=%d ok=%d\n",
         rb.produced, rb.consumed, ok);
}

// ---------------- 测试运行器 ----------------

static void
tests_entry(void)
{
  printf("proc tests: start\n");

  test_process_creation();
  test_scheduler();
  test_synchronization();

  printf("proc tests: all done\n");
  // 作为独立进程结束
  exit_process(0);
}

// 对外暴露：由 main 调用一次即可
void
proc_tests_start(void)
{
  int pid = create_process(tests_entry);
  if (pid < 0) {
    printf("proc_tests_start: create tests_entry failed\n");
  }
}